#
# Copyright (c) 2009, 2010, Davide Cittaro
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the <organization> nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import MSMS
import sys
import numpy as np
# the following are here to read mzXML
import struct
from base64 import b64decode
from xml.dom import minidom

class MSMSAnalysis():
  """ Class to store MS/MS analysis files
  """
  def __init__(self):
    self._spectraList = np.array([])
    self._mzList = np.array([])
    self._massList = np.array([])
    self._searchType = ''
    self._chksumset = set()
  def parseMgfFile(self, mgffile, append = True):
    """ Fills class instance with a provided 
    mascot generic format file. This reads file once in memory and then
    tries to fillup the spectra.
    """
    # don't know why but apparently this is the fastest method
    try:
      fh = open(mgffile, 'r')
    except IOError:
      sys.stderr.write('Missing file ' + mgffile + '\n')
      sys.exit(1)
    n = 0
    if not append:
      self._spectraList = np.array([])

    for line in fh:
      if line[:5] == 'SEARC': continue
      if line[:5] == 'REPTY': continue
      if line[:4] == 'SCAN': continue
      if line[:5] == 'RTINS': continue
      if line[0] == '#': continue
      if not line.strip(): continue
      if line[:5] == 'BEGIN':
        temporarySpectrum = MSMS.ms2spectrum()
        temporaryPeaks = []
      elif line[:7] == 'PEPMASS':
        temporarySpectrum.setPrecursorMz(float(line[8:]))
      elif line[:6] == 'CHARGE':
        temporarySpectrum.setCharge(int(line[7]))
      elif line[:5] == 'TITLE':
        temporarySpectrum.setDescription(line[6:].strip())
      elif line[:3] == 'END':
        temporarySpectrum.setSpectrum(temporaryPeaks)
        chk = temporarySpectrum.checksum()
        if chk not in self._chksumset:
          temporarySpectrum.setFileName(mgffile)
          self._spectraList = np.append(self._spectraList, temporarySpectrum)
          self._chksumset.add(chk)
      else:
        temporaryPeaks.append(line.split())

#    fileLine = fh.read() #we have the file into a line
    fh.close()
#    spLines = fileLine.split('END')
#    if '\r\n' in fileLine:
#      eol = '\r\n'
#    elif '\n' in fileLine:
#      eol = '\n'
#   else:
#      eol = '\r'
#    s = fileLine.find('BEGIN')
#    tempList = fileLine[s:].split('END IONS'+eol) # every element is a spectrum as a line
#    tempList = map(lambda y: y.split(eol), tempList) # every element is a spectrum as list
#    tempList = filter(lambda x: 'BEGIN IONS' in x, tempList)  # remove non spectrum element
#    
#    indxs = range(len(tempList))
#    self._spectraList = np.array(map(lambda x: ms2spectrum(), indxs))  #initialize ms2spectrum() objects
#    map(lambda x: self._spectraList[x].initWithList(tempList[x]), indxs)
    self._mzList = np.array([x.precursorMz() for x in  self._spectraList])
    self._massList = np.array([x.precursorMass() for x in self._spectraList])

#  def parseMzXML(self, mzxmlfile):
#    """ Fills with mzxml file, only msLevel = 2
#    """
#    
#    mzFile = minidom.parse(mzxmlfile)
#    ms2nodes = filter(lambda x: x.attributes["msLevel"].value == '2', mzFile.getElementsByTagName('scan'))
#    
#    for node in ms2nodes:
#      newMSMS = ms2spectrum()
#      
#      newMSMS.setPrecursorMz(float(node.getElementsByTagName('precursorMz')[0].childNodes[0].nodeValue))
#      newMSMS.setCharge(int(node.getElementsByTagName('precursorMz')[0].attributes["precursorCharge"].value))
#      tString = 'ScanNumber: ' + str(node.attributes["num"].value) + 'PrecIntensity: ' + str(node.getElementsByTagName('precursorMz')[0].attributes["precursorIntensity"].value)
#      newMSMS.setDescription(tString)
#      
#      rawPeak = node.getElementsByTagName('peaks')[0].childNodes[0].nodeValue
#
#      masses = []
#      intensities = []
#      decodedString = b64decode(rawPeak)
#      decodeMask = '>' + 'f'*(len(decodedString)/4) 
#      decodedList = struct.unpack(decodeMask, decodedString)
#      
#
#      newMSMS.setSpectrum(map(lambda x: [decodedList[x], decodedList[x+1]], range(0,len(decodedList),2)))
#      
#      self._spectraList.append(newMSMS)
    
  def dumpMgfFile(self, mgffile):
    """ Dumps class content into a mascot generic format file
    """
    fh = open(mgffile, 'w')
    fh.write("SEARCH=MIS\nREPTYPE=PEPTIDE\n")
    for s in self._spectraList:
      fh.write("BEGIN IONS\n")
      fh.write("PEPMASS=" + str(s.precursorMz()) + "\n")
      fh.write("CHARGE=" + str(s.charge()) + "+\n")
      fh.write("TITLE=" + s.description() + "\n")
      for p in s.peaks():
        fh.write(str(p[0, 0]) + " " + str(p[0, 1]) + "\n")
      fh.write("END IONS\n\n")
    fh.close()
#  def fillWithMzxmlFile(self, mzxmlfile):
#    # this will be useful as soon as we understand if the format will be
#    # spread around MS community...
#    fh = open(mzxmlfile, 'r')
#    self._spectraList = []
#    fh.close()
#  def fillWithDtaFile(self, dtafile):
#    fh = open(dtafile, 'r')
#    self._spectraList = []
#    fh.close()
  def mzList(self):
    """ Returns a list of all m/z values for 
    precursors
    """
    if not self._mzList.any():
      return np.array([x.precursorMz() for x in  self._spectraList])
    return self._mzList
  def massList(self):
    """ Returns a list of all Mr values for 
    precursors
    """
    if not self._massList.any():
      return np.array([x.precursorMass() for x in  self._spectraList])
    return self._massList
  def searchType(self):
    """ Returns the search type for MGF files
    """
    return self._searchType
  def  spectra(self):
    """ Return the list of stored spectra objects
    """
    return self._spectraList
  def items(self):
    """ Return an iterator over spectra
    """
    return iter(self._spectraList)
  def length(self):
    """ Return the number of parsed spectra 
    """
    return len(self._spectraList)
  def spectrumAtIndex(self, index = 0):
    """ Return the spectrum at a certain index.
    I think this method is pretty useless
    """
    try:
      return self._spectraList[index]
    except IndexError:
      return None
  def precursorMzSearch(self, mass, tolerance = 0.5, unit = 'Da'):
    """ Return a new MSMSAnalysis object with all the spectra
    with a precursor m/z equal to the mass provided, within
    a certain tolerace
    """
    t = MSMSAnalysis()
    mass = float(mass)
    tolerance = float(tolerance)

    if unit == 'Da':
      None
    elif unit == 'ppm':
      tolerance = mass / 1000000 * tolerance
    else:
      return None
      
    t._spectraList = np.array([x for x in  self._spectraList if abs(x.precursorMz() - mass) <= tolerance])
    t._mzList = t.mzList()
    t._massList = t.massList()
    return t
  def precursorMassSearch(self, mass, tolerance = 0.5, unit = 'Da'):
    """ Return a new MSMSAnalysis object with all the spectra
    with a precursor Mr equal to the mass provided, within
    a certain tolerace
    """
    t = MSMSAnalysis()
    mass = float(mass)
    tolerance = float(tolerance)

    if unit == 'Da':
      None
    elif unit == 'ppm':
      tolerance = mass / 1000000 * tolerance
    else:
      return None

    t._spectraList = np.array([x for x in self._spectraList if abs(x.precursorMass() - mass) <= tolerance])
    t._mzList = t.mzList()
    t._massList = t.massList()
    return t
  def chargeSearch(self, charge):
    """ Return a new MSMSAnalysis object with all the spectra
    with a certain charge value
    """
    t = MSMSAnalysis()
    charge = int(charge)
    t._spectraList = np.array([x for x in self._spectraList if x._charge == charge])
    t._mzList = t.mzList()
    t._massList = t.massList()
    return t
  def descriptionSearch(self, string):
    """ Return a new MSMSAnalysis object with all the spectra
    with given text into spectrum description. Useful for filtering
    selected spectra...
    """
    t = MSMSAnalysis()
    t._spectraList = np.array([x for x in self._spectraList if string in x._description])
    t._mzList = t.mzList()
    t._massList = t.massList()
    return t
  def addSpectrum(self, spectrum):
    """ Add a new spectrum to this analysis
    """
    self._spectraList = np.append(self._spectraList, spectrum)
  def reduceNoise(self, ratio):
    """ Cut every peak that has intensity lower that a
    certain ratio to the max intensity peak
    """
    [x.reduceNoise(ratio) for x in  self.spectra()]
  def topIons(self, nIons = 12, window = 30):
    """ select the top nIons in window for each spectrum
    this should enhance specificity of the method...
    """
    self._spectraList = np.array([x.topIons(nIons, window) for x in  self._spectraList])

